package com.sc.nft.config.aspect;

import cn.hutool.core.lang.Assert;
import com.sc.nft.annotation.PartitionField;
import com.sc.nft.annotation.PartitionPK;
import com.sc.nft.annotation.PartitionTable;
import com.sc.nft.helper.TablePartitionRouter;
import com.sc.nft.sup.EntityContract;
import com.sc.nft.util.TableNameContext;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;

import java.lang.annotation.Annotation;
import java.lang.reflect.Method;

import static com.sc.nft.constant.TableNameConstants.*;

/**
 * MySQL表分片切面配置
 *
 * @author zlikun
 * @since 2024/06/06
 */
@Slf4j
@Aspect
@Component
public class TablePartitionAspect {

    private final String PARTITION_ARG_NAME_ID = "id";
    private final String PARTITION_ARG_NAME_FIELD = "field";

    @Around("@annotation(pt)")
    public Object around(ProceedingJoinPoint pjp, PartitionTable pt) throws Throwable {

        Assert.notBlank(pt.value(), "@PartitionTable 注解值不能为空");

        Object[] args = pjp.getArgs();
        Assert.notEmpty(args, "@PartitionTable 注解修饰方法必须包含分区字段参数");

        // 获取方法分区参数注解
        MethodSignature signature = (MethodSignature) pjp.getSignature();
        Method method = signature.getMethod();

        // 获取分区参数注解，两者并存时优先主键，其次分区键
        String partitionArgName = null;
        Object partitionArgValue = null;

        Annotation[][] annotations = method.getParameterAnnotations();
        label:
        for (int i = 0; i < annotations.length; i++) {
            for (Annotation an : annotations[i]) {
                if (an instanceof PartitionPK) {
                    partitionArgName = PARTITION_ARG_NAME_ID;
                    if (args[i] instanceof EntityContract) {
                        partitionArgValue = ((EntityContract) args[i]).getId();
                    } else {
                        partitionArgValue = args[i];
                    }
                    break label;
                }
                if (an instanceof PartitionField) {
                    partitionArgName = PARTITION_ARG_NAME_FIELD;
                    partitionArgValue = args[i];
                    break label;
                }
            }
        }

        Assert.notNull(partitionArgValue, "@PartitionTable 注解修饰方法必须指定分区字段参数，且参数值不能为空");

        String partition = null;
        switch (pt.value()) {
            case COLLECTION_CHAIN_INIT_DETAILS_LOG:
                partition = defaultPartitionMethod(TablePartitionRouter.ofCollectionChainInitDetailsLogRouter(), partitionArgName, partitionArgValue);
                break;
            case EXPENDITURE_RECORD:
                partition = defaultPartitionMethod(TablePartitionRouter.ofExpenditureRecord(), partitionArgName, partitionArgValue);
                break;
            case USER_WALLET_DETAIL:
                partition = defaultPartitionMethod(TablePartitionRouter.ofUserWalletDetail(), partitionArgName, partitionArgValue);
                break;
            default:
                break;
        }

        // 不分片
        if (partition == null) {
            return pjp.proceed();
        }

        log.debug("MySQL表执行分片机制 table: {}, partition: {}", pt.value(), partitionArgName);

        try {
            TableNameContext.setPartition(partition);
            return pjp.proceed();
        } catch (Throwable e) {
            log.error("MySQL表{}分片切面执行出错 partition: {}", pt.value(), partition, e);
            throw new RuntimeException(e);
        } finally {
            TableNameContext.clear();
        }
    }

    private String defaultPartitionMethod(TablePartitionRouter.Router router, String partitionArgName, Object partitionArgValue) {
        if (!(partitionArgValue instanceof Number)) {
            return null;
        }
        switch (partitionArgName) {
            case PARTITION_ARG_NAME_FIELD:
                return router.routeWithField(((Number) partitionArgValue));
            case PARTITION_ARG_NAME_ID:
                return router.routeWithPk(((Number) partitionArgValue));
            default:
                throw new IllegalArgumentException("不支持的分区字段");
        }
    }

}
